import sys
from PyQt5.QtWidgets import QApplication, QWidget, QGridLayout, QLabel 
from PyQt5.QtGui import QPixmap
import sys
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
import SimpleITK as sitk
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

begin_select = False

label_color = {
    1: [249,72,247],
    2: [209,198,103],
    3: [70,139,88],
    4: [199,104,19],
    5: [174,213,76],
    6: [83,83,50]
}

class Display(QWidget):
    def __init__(self, image_arr, label_arr):
        super().__init__()
        self.image_arr = image_arr
        self.label_arr = label_arr
        self.initUI()
        
    def initUI(self):
        # 创建一个网格布局
        self.grid = QGridLayout()
         
        # area [0, 0]
        self.grid.addLayout(Grid(self.image_arr, self.label_arr,  0), 1, 0)
        
        # area [0, 1]
        self.grid.addLayout(Grid(self.image_arr, self.label_arr, 1), 0, 0)
        
        # area [1, 0]
        self.grid.addLayout(Grid(self.image_arr, self.label_arr, 2), 0, 1)
        
        self.btn = QPushButton()
        self.btn.setCheckable(True)
        self.btn.setText("begin select point")
        self.btn.clicked.connect(self.update_status)
        self.grid.addWidget(self.btn, 1, 1)
        
        # 设置窗口的布局为网格布局
        self.setLayout(self.grid)
 
        # 设置窗口的大小
        self.setGeometry(300, 300, 400, 300)
        self.setWindowTitle('Image Grid')
        self.show()
    
    def update_status(self):
        global begin_select
        if self.btn.isChecked():
            begin_select = True
        else:
            begin_select = False
 
 
class Grid(QHBoxLayout):
    def __init__(self, image_arr, label_arr, type):
        super().__init__()
        self.image_arr = image_arr.copy()
        self.label_arr = label_arr.copy()
        self.type = type
        self.initUI()
 
    def initUI(self):
        width, height, channel = self.image_arr.shape                 
        self.imagewidget = ImageWidget()
        if self.type == 0:
            image_fusion = self.img_fusion(self.image_arr[int((width - 1)/2), :, :], self.label_arr[int((width - 1)/2), :, :])
            self.imagewidget.load_image(image_fusion, int((width - 1)/2), 0)
            self.scroll = QScrollBar()
            self.scroll.setMaximum(width)
            self.scroll.setValue(int((width - 1) / 2))
        elif self.type == 1:
            image_fusion = self.img_fusion(self.image_arr[:, int((height - 1)/2), :], self.label_arr[:, int((height - 1)/2), :])
            self.imagewidget.load_image(image_fusion, int((height - 1)/2), 0)
            self.scroll = QScrollBar()
            self.scroll.setMaximum(height)
            self.scroll.setValue(int((height - 1) / 2))
        else:
            image_fusion = self.img_fusion(self.image_arr[:, :, int((channel - 1)/2)], self.label_arr[:, :, int((channel - 1)/2)])
            self.imagewidget.load_image(image_fusion, int((channel - 1)/2), 0)
            self.scroll = QScrollBar()
            self.scroll.setMaximum(channel)
            self.scroll.setValue(int((channel - 1) / 2))
        
        self.scroll.sliderMoved.connect(self.sliderval)
        self.addWidget(self.scroll)
        self.addWidget(self.imagewidget)
    
    def sliderval(self):
        print( self.scroll.value())
        if self.type == 0:
            image_fusion = self.img_fusion(self.image_arr[self.scroll.value() -1, :, :], self.label_arr[self.scroll.value() -1, :, :])
            self.imagewidget.load_image(image_fusion, self.scroll.value() -1, 0)
        elif self.type == 1:
            image_fusion = self.img_fusion(self.image_arr[:, self.scroll.value() - 1, :], self.label_arr[:, self.scroll.value() - 1, :])
            self.imagewidget.load_image(image_fusion, self.scroll.value() - 1, 0)
        else:
            image_fusion = self.img_fusion(self.image_arr[:, :, self.scroll.value() -1], self.label_arr[:, :, self.scroll.value() -1])
            self.imagewidget.load_image(image_fusion, self.scroll.value() -1, 0)
    def img_fusion(self, image, label):
        image_fusion = np.expand_dims(image, axis=-1) 
        image_fusion = np.repeat(image_fusion, 3, axis=-1)
        width = label.shape[1]
        height = label.shape[0]
        for i in range(width):
            for j in range(height):
                if label[j,i] == 0:
                    continue
                
                image_fusion[j,i] = label_color[label[j,i]]

        return image_fusion

    
class ImageWidget(QWidget):
    def __init__(self, parent=None):
        super(ImageWidget, self).__init__(parent)
        self.image = None
        self.image_loaded = False
        self.index = 0
        self.type = 0
 
    def load_image(self, image, index, type):
        self.image = image.copy()
        self.image_loaded = True
        self.update()  # Trigger repaint
        self.index = index
        self.type = type
        
 
    def paintEvent(self, event):
        if self.image_loaded:
            painter = QPainter(self)
            image_rgb = self.image
            # image_rgb = np.expand_dims(self.image, axis=-1) 
            # image_rgb = np.repeat(image_rgb, 3, axis=-1)
            
            width = self.image.shape[1]
            height = self.image.shape[0]
            self.setFixedSize(QSize(width, height))
            format = QImage.Format_RGB888

            qimage = QImage(image_rgb.data, width, height, width * 3, format) #
            painter.drawImage(self.rect(), qimage)
                    
    
    def mousePressEvent(self, event):
        global begin_select
        if event.button() == Qt.LeftButton and begin_select:
            
            # 获取鼠标点击位置
            x = event.pos().x()
            y = event.pos().y()
            
            local_x = 0
            local_y = 0
            local_z = 0  
            if self.type == 0:
                local_x = self.index
                local_y = x
                local_z = y
            elif self.type == 1:
                local_x = x
                local_y = self.index
                local_z = y
            else:
                local_x = x
                local_y = y
                local_z = self.index
            
            # 显示鼠标点击位置            
            status = QMessageBox.question(self,"Whether to start SAM inference?",
                                          f"srceen postion:[{x}, {y}]\n local postion:[{local_x},{local_y}, {local_z}]",
                                          QMessageBox.No | QMessageBox.Yes)#创建一个二次确认框
            
            if status == QMessageBox.Yes:
                print(f"width:{self.width()}, height: {self.height()}")
                print('begin inference'+'\n')#打印当前时间到UI显示，提示完成重启
            else:
                print('reselect point'+'\n')#打印当前时间到UI显示，提示取消重启
 
if __name__ == '__main__':

    image = sitk.ReadImage("/media/chyang/data/dataset/project/hipbone/nnUNet_raw/Dataset001_PelvicNAS/imagesTr/1.3.6.1.4.1.9328.50.4.0001_0000.nii.gz") #in_file是nii.gz文件的路径
    image_arr = sitk.GetArrayFromImage(image)
    img_scaled = (image_arr - image_arr.min()) / (image_arr.max() - image_arr.min()) * 255
    img_scaled = img_scaled.astype(np.uint8)

    label = sitk.ReadImage("/media/chyang/data/dataset/project/hipbone/nnUNet_raw/Dataset001_PelvicNAS/labelsTr/1.3.6.1.4.1.9328.50.4.0001.nii.gz")
    label_arr = sitk.GetArrayFromImage(label)
    label_scaled = label_arr
    label_scaled = label_scaled.astype(np.uint8)

    
    app = QApplication(sys.argv)
    ex = Display(img_scaled, label_scaled)
    sys.exit(app.exec_())
    